#pragma comment(lib, "d3d9.lib")
#pragma comment(lib, "d3dx9.lib")

#include <Windows.h>
#include <stdio.h>
#include <fstream>
#include <iostream>
#include <d3d9.h>
#include <d3dx9.h>

#define CREATE_DEVICE 16 	//vTable address of the createdevice function 
#define ENDSCENE 42			//vTable address of the endscene function 

///////////////////////////////////////////////////////
//init stuff-typedefs etc

//be sure to get the right function params check d3dx9.h for details or vTable addresses
//createdevice functionPtr
typedef HRESULT (WINAPI *CreateDevice_t)(IDirect3D9* Direct3D_Object, UINT Adapter, D3DDEVTYPE DeviceType, HWND hFocusWindow, 
                    DWORD BehaviorFlags, D3DPRESENT_PARAMETERS* pPresentationParameters, 
                    IDirect3DDevice9** ppReturnedDeviceInterface);

//endscene functionPtr
typedef HRESULT (WINAPI *EndScene_t)(IDirect3DDevice9* surface);


// Function pointers to the original functions
CreateDevice_t D3DCreateDevice_orig;
EndScene_t D3DEndScene_orig, D3DEndScene_trampoline;

//this is just some memory for our trampoline function
struct trampoline{
	char code[5], jmp[5];
};


// Our hooking functions-these are used for initial hooking
HRESULT WINAPI D3DCreateDevice_hook(IDirect3D9* Direct3D_Object, UINT Adapter, D3DDEVTYPE DeviceType, HWND hFocusWindow, 
                    DWORD BehaviorFlags, D3DPRESENT_PARAMETERS* pPresentationParameters, 
                    IDirect3DDevice9** ppReturnedDeviceInterface);

HRESULT WINAPI D3DEndScene_hook(IDirect3DDevice9* device);

////////////////////////////////////////////////////////////////////////
//prototypes
HRESULT HookCreateDevice();
DWORD WINAPI vTableThread(LPVOID);
DWORD WINAPI DX_Init(LPVOID);
////////////////////////////////////////////////////////////////////////
//variables
PDWORD vTable = NULL;
FILE *file;
LPD3DXFONT g_Font = NULL;
LPDIRECT3DDEVICE9 device = NULL;
RECT FontPosition;


//////////////////////////////////////////////////////////////////////
//implementation

//DLL Entry (shouldn't be a surprise
BOOL APIENTRY DllMain(HMODULE hModule, DWORD  ul_reason_for_call, LPVOID lpReserved){
	switch(ul_reason_for_call){
	case DLL_PROCESS_ATTACH:
		//just to mark the position where our hook routine is
		fopen_s(&file, "tmp.txt", "a+");
		fprintf_s(file, "----------------------Process Attached-----------------------------\n");
		
		//this would hook the app in the beginning(must be started before d3dcreatedevice is called
		//HookCreateDevice();
		
		//Starting our thread for the magic
		CreateThread(NULL, NULL, DX_Init, NULL, NULL, NULL);
		break;
	case DLL_PROCESS_DETACH:
		fprintf_s(file, "----------------------Process Detached-----------------------------\n");
		fclose(file);
		UnmapViewOfFile(shared);
		CloseHandle(map);
		break;
	}
	return TRUE;
}

//hooking the createdevice function to create our own device, the proc uses then our d3dxdevice to paint it's stuff-good for us
HRESULT HookCreateDevice(){
	fprintf_s(file, "setting the CreateDevice-Hook\n");
	IDirect3D9 *device;
	//calculate the function address
	device = Direct3DCreate9(D3D_SDK_VERSION);
	vTable = (DWORD*)*(DWORD*)device;
	device->Release();
	DWORD protectFlag;
	//setting write permissions
	if(VirtualProtect(&vTable[CREATE_DEVICE], sizeof(DWORD), PAGE_READWRITE, &protectFlag)){
		*(DWORD*)&D3DCreateDevice_orig = vTable[CREATE_DEVICE];
		*(DWORD*)&vTable[CREATE_DEVICE] = (DWORD)D3DCreateDevice_hook;
		if(!VirtualProtect(&vTable[CREATE_DEVICE], sizeof(DWORD), protectFlag, &protectFlag)){
			MessageBox(NULL, "failed to hook CreateDevice\n", "D3D9CreateDeviceHook", MB_ICONEXCLAMATION);
			return D3DERR_INVALIDCALL;
		}
	}
	else
		return D3DERR_INVALIDCALL;
	fprintf_s(file, "CreateDevice-Hook OK\n
	//thats i when the app now wants to create a device our routine is called --> D3DCreateDevice_hook
	return D3D_OK;
}


//here we just create our normal device but we store it's endscene function to insert our code there
HRESULT WINAPI D3DCreateDevice_hook(IDirect3D9* Direct3D_Object, UINT Adapter, D3DDEVTYPE DeviceType, HWND hFocusWindow, 
                    DWORD BehaviorFlags, D3DPRESENT_PARAMETERS* pPresentationParameters, 
                    IDirect3DDevice9** ppReturnedDeviceInterface)
{
	fprintf_s(file, "CreateDevice called, setting up the device...\n");
	HRESULT result = D3DCreateDevice_orig(Direct3D_Object, Adapter, DeviceType, hFocusWindow, BehaviorFlags|D3DCREATE_MULTITHREADED, pPresentationParameters,
		ppReturnedDeviceInterface);
	DWORD protectFlag;
	if(VirtualProtect(&vTable[CREATE_DEVICE], sizeof(DWORD), PAGE_READWRITE, &protectFlag)){
		fprintf_s(file, "resetting the CreateDevice function\n");
		*(DWORD*)&vTable[CREATE_DEVICE] = (DWORD)D3DCreateDevice_orig;
		if(!VirtualProtect(&vTable[CREATE_DEVICE], sizeof(DWORD), protectFlag, &protectFlag)){
			fprintf_s(file, "error while resetting function\n");
			return D3DERR_INVALIDCALL;
		}
	}
	else{
		fprintf_s(file, "error while writing memory\n");
		return D3DERR_INVALIDCALL;
	}
	if(result == D3D_OK){
		fprintf_s(file, "successfully hooked endscene-function. Creating vTable-thread\nvTable: %x\n", vTable);
		vTable = (DWORD*)*(DWORD*)*ppReturnedDeviceInterface;
		*(PDWORD)&D3DEndScene_orig = (DWORD)vTable[ENDSCENE];
		CreateThread(NULL, NULL, vTableThread, NULL, NULL, NULL);
	}
	fprintf_s(file, "test: %d (%x)\n", &vTable[ENDSCENE], &vTable[ENDSCENE]);
	return result;
}


DWORD WINAPI vTableThread(LPVOID args){
	while(TRUE){
		Sleep(100);
		*(PDWORD)&vTable[ENDSCENE] = (DWORD)D3DEndScene_hook;
	}
}

ID3DXFont *g_font=NULL;
void initFont();
bool run = false;


//that's the endscene hook *magic*
HRESULT WINAPI D3DEndScene_hook(IDirect3DDevice9* dev){
	fprintf_s(file, "EndScene called\n");
	//setting our font
	initFont(IDirect3DDevice9* dev);
	FontPosition.top = 100;         //position
	g_Font->DrawText(NULL,"Direct3D9 Text With D3DXFont",-1,&FontPosition,DT_CENTER,0xffffffff); //draw text
	//and here our trampoline gets called(original endscene function)
	HRESULT result = D3DEndScene_trampoline(dev);
	return result;
}


LRESULT CALLBACK MsgProc(HWND hwnd,UINT uMsg,WPARAM wParam,LPARAM lParam){
	return DefWindowProc(hwnd, uMsg, wParam, lParam);
}


bool bCompare(const BYTE* pData, const BYTE* bMask, const char* szMask)
{
    for(;*szMask;++szMask,++pData,++bMask)
        if(*szMask=='x' && *pData!=*bMask)   return 0;
    return (*szMask) == NULL;
}

DWORD FindPattern(DWORD dwdwAdd,DWORD dwLen,BYTE *bMask,char * szMask)
{
    for(DWORD i=0; i<dwLen; i++)
        if (bCompare((BYTE*)(dwdwAdd+i),bMask,szMask))  return (DWORD)(dwdwAdd+i);
    return 0;
}


void initFont(IDirect3DDevice9* dev){
	if(g_Font != NULL)
		g_Font->Release();
	D3DXFONT_DESC FontDesc = {24,
                          0,
                          400,
                          0,
                          false,
                          DEFAULT_CHARSET,
                          OUT_TT_PRECIS,
                          CLIP_DEFAULT_PRECIS,
                          DEFAULT_PITCH,
                          "Arial"};
	FontPosition.top = 0;
	FontPosition.left = 0;
	FontPosition.right = 400;
	FontPosition.bottom = 400;
	D3DXCreateFontIndirectA(dev,&FontDesc,&g_Font);
}


//this is the hook routine called by our thread and where the magic is done!
//note that this routine can be much more complicated since this also works when the proc already created its d3dx device 
//detailed asm knowledge is necessary!
DWORD WINAPI DX_Init(LPVOID args)
{
	//first create a window because per window only one d3dx device may exist (just to get the endscene function address)
	fprintf_s(file, "Initializing (DX_INIT)\n");
	PDWORD t;
	WNDCLASSEXA wc = {sizeof(WNDCLASSEX), CS_CLASSDC, MsgProc, 0L, 0L, GetModuleHandleA(NULL), NULL, NULL, NULL, NULL, "DX", NULL};
    RegisterClassExA(&wc);
    HWND hWnd = CreateWindowA("DX",NULL,WS_OVERLAPPEDWINDOW,100,100,300,300,GetDesktopWindow(),NULL,wc.hInstance,NULL);
	//Normal d3dx setup
	fprintf_s(file, "created window, creating dummy d3d9device\n");
    LPDIRECT3D9 pD3D = Direct3DCreate9( D3D_SDK_VERSION );
	fprintf_s(file, "init LPDIRECT3D9 //OK\n");
    D3DPRESENT_PARAMETERS d3dpp;
    ZeroMemory( &d3dpp, sizeof(d3dpp) );
    d3dpp.Windowed = TRUE;
    d3dpp.SwapEffect = D3DSWAPEFFECT_DISCARD;
    d3dpp.BackBufferFormat = D3DFMT_UNKNOWN;
	fprintf_s(file, "init D3DPRESENT_PARAMETERS //OK\n");
    LPDIRECT3DDEVICE9 pd3dDevice;
	pD3D->CreateDevice(D3DADAPTER_DEFAULT,D3DDEVTYPE_NULLREF,hWnd,D3DCREATE_SOFTWARE_VERTEXPROCESSING,&d3dpp,&pd3dDevice);
	fprintf_s(file, "init LPDIRECT3DDEVICE9 //OK\n");
	t = (PDWORD)*(PDWORD)pd3dDevice;
	DWORD oldFlag;
//here the endscene function is hooked
	struct trampoline *tr = (struct trampoline*)calloc(1, sizeof(struct trampoline));
	if(VirtualProtect((PDWORD)t[ENDSCENE], sizeof(DWORD)*2, PAGE_READWRITE, &oldFlag)){
//because we'll gonna steel 5bytes of the original function we will first save them in our memory struct exec_readwrite because we want to exec it later
		VirtualProtect(tr, sizeof(struct trampoline), PAGE_EXECUTE_READWRITE, &nf);
//and there it is, our trampoline
		*(PDWORD)&D3DEndScene_trampoline = (DWORD)tr;
		for(int i = 0; i < 5; i++)
			tr->code[i] = ((char*)t[ENDSCENE])[i];
//first of all store the original entry point of the function (-->endscene address now in D3DEndScene_orig)
		*(PDWORD)&D3DEndScene_orig = t[ENDSCENE];
//now change the opcode at the beginning to unconditional jmp
		*(char*)t[ENDSCENE] = 0xe9;
//then infuse the jmp address where our endscene function is (-->relative jmp means jmp address = to-from-5, the 5 comes from the opcode length jmp takes 5bytes)
		*(PDWORD)&((char*)t[ENDSCENE])[1] = (DWORD)D3DEndScene_hook - (DWORD)t[ENDSCENE] - 5;
//this is the jmp to the orig function from our trampoline (to the continuing part)
		tr->jmp[0] = 0xe9;
		*(PDWORD)&(tr->jmp[1]) = (DWORD)D3DEndScene_orig - (DWORD)&tr->jmp[0];
		if(!VirtualProtect((PDWORD)t[ENDSCENE], sizeof(DWORD)*2, oldFlag, &oldFlag)){
			fprintf_s(file, "could not override endscene\n");
			MessageBox(NULL, "Could not ovveride EndScene!", "d3dHook", MB_ICONERROR);
		}
		else
			fprintf_s(file, "overrided EndScene Method\n");
		
	}
//cleanup and that's it
    DestroyWindow(hWnd);
	fprintf_s(file, "finished DX_INIT\n");
	return 0;
}